"""
    解码器及解码层
"""
import tensorflow as tf
from .multi_head_attention import MultiHeadAttention
from .common_function import point_wise_feed_forward_network, positional_encoding
import config


class DecoderLayer(tf.keras.layers.Layer):
    def __init__(self):
        """  解码层
            参数:
                1.d_model    --> 字向量维度(embedding_size)
                2.num_heads  --> 多头注意力头的个数
                3.dff        --> 全连接神经元数量
                4.rate       --> drop层几率
        """
        super(DecoderLayer, self).__init__()
        # 初始化参数
        self.d_model = config.d_model
        self.num_heads = config.num_heads
        self.dff = config.dff
        self.rate = config.dropout_rate


        self.mha1 = MultiHeadAttention(self.d_model, self.num_heads)
        self.mha2 = MultiHeadAttention(self.d_model, self.num_heads)

        self.ffn = point_wise_feed_forward_network(self.d_model, self.dff)

        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = tf.keras.layers.Dropout(self.rate)
        self.dropout2 = tf.keras.layers.Dropout(self.rate)
        self.dropout3 = tf.keras.layers.Dropout(self.rate)

    def call(self, x, enc_output, training,
             look_ahead_mask, padding_mask):
        # enc_output.shape == (batch_size, input_seq_len, d_model)

        attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)
        attn1 = self.dropout1(attn1, training=training)
        out1 = self.layernorm1(attn1 + x)

        attn2, attn_weights_block2 = self.mha2(
            enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)
        attn2 = self.dropout2(attn2, training=training)
        out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)

        ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)
        ffn_output = self.dropout3(ffn_output, training=training)
        out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)

        return out3, attn_weights_block1, attn_weights_block2


# 解码器
class Decoder(tf.keras.layers.Layer):
    def __init__(self, target_vocab_size,
                 maximum_position_encoding):
        """ 解码器
             参数
                 1.num_layers --> 编码层数量
                 2.d_model    -->字向量维度(embedding_size)
                 3.num_heads  -->多头注意力头的个数
                 4.dff        -->全连接神经元数量
                 5.target_vocab_size -->输出字向量的大小(字典大小)
                 6.maximun_position_encoding --> 输出位置编码信息
                 7.rate       -->drop层几率
        """
        super(Decoder, self).__init__()
        # 解码器隐藏层
        self.d_model = config.d_model
        # 编码层数量
        self.num_layers = config.num_layers

        self.rate = config.dropout_rate

        self.embedding = tf.keras.layers.Embedding(target_vocab_size, self.d_model)
        self.pos_encoding = positional_encoding(maximum_position_encoding, self.d_model)

        self.dec_layers = [DecoderLayer()
                           for _ in range(self.num_layers)]
        self.dropout = tf.keras.layers.Dropout(self.rate)

    def call(self, x, enc_output, training,
             look_ahead_mask, padding_mask):
        seq_len = tf.shape(x)[1]
        attention_weights = {}

        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x, block1, block2 = self.dec_layers[i](x, enc_output, training,
                                                   look_ahead_mask, padding_mask)

            attention_weights['decoder_layer{}_block1'.format(i + 1)] = block1
            attention_weights['decoder_layer{}_block2'.format(i + 1)] = block2

        # x.shape == (batch_size, target_seq_len, d_model)
        return x, attention_weights
